# -*- coding: utf-8 -*-
"""
RandomForest股票预测模型训练脚本
"""

import sys
import os
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error
import matplotlib.pyplot as plt
import joblib
import argparse
from pypinyin import lazy_pinyin

def chinese_to_pinyin(text):
    """将中文转换为拼音"""
    if any('\u4e00' <= char <= '\u9fff' for char in text):
        return ' '.join(lazy_pinyin(text)).title()
    return text

def train_randomforest_model(stock_code, frequency):
    """训练RandomForest模型"""
    
    print(f"🌳 开始训练RandomForest模型 - 股票: {stock_code}, 周期: {frequency}")
    print("=" * 60)
    
    # 文件名
    input_filename = f'output/data/stock_data_{stock_code}_{frequency}_with_indicators.csv'
    model_filename = f'output/models/rf_model_{stock_code}_{frequency}.pkl'
    plot_filename = f'output/plots/prediction_vs_actual_{stock_code}_{frequency}_randomforest.png'
    
    try:
        # 加载数据
        print("📊 加载数据...")
        df = pd.read_csv(input_filename)
        
        # 修复列名问题
        if 'date' in df.columns:
            df.rename(columns={'date': '日期'}, inplace=True)
        
        # 设置日期索引
        df['日期'] = pd.to_datetime(df['日期'])
        df.set_index('日期', inplace=True)
        
        print(f"   数据规模: {len(df)} 条记录")
        
        # 准备特征和目标变量
        feature_columns = [col for col in df.columns if col not in ['code', 'close']]
        target_column = 'close'
        
        # 创建目标变量（预测下一天收盘价）
        df['target'] = df[target_column].shift(-1)
        df = df.dropna()
        
        # 准备训练数据
        X = df[feature_columns]
        y = df['target']
        
        print(f"   特征数量: {len(feature_columns)}")
        print(f"   有效样本: {len(X)}")
        
        # 分割数据（按时间序列，前80%训练，后20%测试）
        split_idx = int(0.8 * len(df))
        X_train, X_test = X[:split_idx], X[split_idx:]
        y_train, y_test = y[:split_idx], y[split_idx:]
        
        print(f"   训练集大小: {len(X_train)}")
        print(f"   测试集大小: {len(X_test)}")
        
        # 训练RandomForest模型
        print("\\n🌳 训练RandomForest模型...")
        model = RandomForestRegressor(
            n_estimators=200,      # 树的数量
            max_depth=10,          # 最大深度  
            min_samples_split=5,   # 分裂最小样本数
            min_samples_leaf=2,    # 叶子节点最小样本数
            random_state=42,       # 随机种子
            n_jobs=-1,            # 使用所有CPU核心
            oob_score=True        # 计算袋外分数
        )
        
        model.fit(X_train, y_train)
        
        # 预测
        print("📊 模型预测...")
        y_pred = model.predict(X_test)
        
        # 计算指标
        mse = mean_squared_error(y_test, y_pred)
        mae = mean_absolute_error(y_test, y_pred)
        rmse = np.sqrt(mse)
        
        # 计算准确率指标
        mape = np.mean(np.abs((y_test - y_pred) / y_test)) * 100
        r2_score = model.score(X_test, y_test)
        oob_score = model.oob_score_
        
        print("\\n📈 模型性能指标:")
        print(f"   RMSE (均方根误差): {rmse:.4f}")
        print(f"   MAE (平均绝对误差): {mae:.4f}")
        print(f"   MAPE (平均绝对百分比误差): {mape:.2f}%")
        print(f"   R² Score: {r2_score:.4f}")
        print(f"   OOB Score (袋外分数): {oob_score:.4f}")
        
        # 保存模型
        print(f"\\n💾 保存模型: {model_filename}")
        # 确保output/models目录存在
        os.makedirs('output/models', exist_ok=True)
        joblib.dump(model, model_filename)
        
        # 特征重要性分析
        feature_importance = pd.DataFrame({
            'feature': feature_columns,
            'importance': model.feature_importances_
        }).sort_values('importance', ascending=False)
        
        print("\\n🎯 前10重要特征:")
        for i, (_, row) in enumerate(feature_importance.head(10).iterrows()):
            print(f"   {i+1:2d}. {row['feature']:<15} {row['importance']:.4f}")
        
        # 创建预测对比图
        print(f"\\n📊 生成预测对比图: {plot_filename}")
        create_prediction_plot(y_test, y_pred, stock_code, frequency, plot_filename)
        
        # 最新预测
        print("\\n🔮 最新预测:")
        latest_data = X.iloc[-1:].values
        next_pred = model.predict(latest_data)[0]
        current_price = df['close'].iloc[-1]
        change_amount = next_pred - current_price
        change_percent = (change_amount / current_price) * 100
        
        print(f"   当前收盘价: {current_price:.2f}元")
        print(f"   预测收盘价: {next_pred:.2f}元")
        print(f"   预测涨跌: {change_amount:+.2f}元 ({change_percent:+.2f}%)")
        
        trend = "📈 看涨" if change_percent > 0 else "📉 看跌" if change_percent < 0 else "➡️ 持平"
        print(f"   预测趋势: {trend}")
        
        print("\\n✅ RandomForest模型训练完成!")
        
    except Exception as e:
        print(f"❌ 训练失败: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

def create_prediction_plot(y_true, y_pred, stock_code, frequency, filename):
    """创建预测对比图"""
    
    # 设置matplotlib
    plt.rcParams['font.family'] = ['DejaVu Sans']
    
    # 只显示最后200个点
    n_show = min(200, len(y_true))
    y_true_show = y_true[-n_show:]
    y_pred_show = y_pred[-n_show:]
    
    # 创建图表
    plt.figure(figsize=(12, 8))
    
    # 主图：预测对比
    plt.subplot(2, 1, 1)
    plt.plot(range(n_show), y_true_show, label=chinese_to_pinyin('实际价格'), 
             color='blue', alpha=0.8, linewidth=1.5)
    plt.plot(range(n_show), y_pred_show, label=chinese_to_pinyin('预测价格'), 
             color='red', alpha=0.8, linewidth=1.5, linestyle='--')
    
    plt.title(f'{stock_code} - RandomForest {chinese_to_pinyin("价格预测对比")} ({frequency})', 
              fontsize=14, fontweight='bold')
    plt.xlabel(chinese_to_pinyin('时间点'))
    plt.ylabel(chinese_to_pinyin('价格') + ' (yuan)')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # 子图：预测误差
    plt.subplot(2, 1, 2)
    errors = y_pred_show - y_true_show
    plt.plot(range(n_show), errors, label=chinese_to_pinyin('预测误差'), 
             color='green', alpha=0.7, linewidth=1)
    plt.axhline(y=0, color='black', linestyle='-', alpha=0.5)
    
    plt.title(f'{stock_code} - RandomForest {chinese_to_pinyin("预测误差分析")}', 
              fontsize=12)
    plt.xlabel(chinese_to_pinyin('时间点'))
    plt.ylabel(chinese_to_pinyin('误差') + ' (yuan)')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # 添加统计信息
    rmse = np.sqrt(np.mean(errors**2))
    mae = np.mean(np.abs(errors))
    plt.text(0.02, 0.98, f'RMSE: {rmse:.2f}\\nMAE: {mae:.2f}', 
             transform=plt.gca().transAxes, verticalalignment='top',
             bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
    
    plt.tight_layout()
    # 确保output/plots目录存在
    os.makedirs('output/plots', exist_ok=True)
    plt.savefig(filename, dpi=300, bbox_inches='tight', facecolor='white')
    plt.close()

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='RandomForest股票预测模型训练')
    parser.add_argument('stock_code', help='股票代码')
    parser.add_argument('frequency', help='数据频率 (daily/hourly)')
    
    args = parser.parse_args()
    
    train_randomforest_model(args.stock_code, args.frequency)

if __name__ == "__main__":
    main()